#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""基于栅格的分区统计
分区统计工具用于根据来自其他数据集的值（赋值栅格）为每一个由区域数据集定义的区域计算统计数据。 为输入区域数据集中的每一个区域计算单个输出值。

通过提供分区图层(矢量或栅格均可)和值图层, 计算出对应区域的统计值, 包括: 平均数、均值、中位数、标准差、方差等。
输出结果能够以 csv 或 json 格式文件保存, 通过 -f [csv|json] 指定。
默认的文件名为 out.[csv|json] , 也可以通过 -o out_filename 指定。

Example:
    $ python calculate-zonal-statistics.py -f csv zonal.tif raster.tif
    $ python calculate-zonal-statistics.py -f csv zonal.shp raster.tif
"""

import os
import sys
import json
import csv
import numpy as np
from osgeo import gdal, osr, ogr

def zonal_stats(feature, input_shapefile, input_raster):
    """分别统计 shapefile 中的每个多边形与值图层叠加后的统计值

    参数:
        feature (int): 图层编号
        input_shapefile (str): 输入 shapefile 文件路径
        input_raster (str): 输入 tif 文件路径

    返回:
        (dict) 单独多边形与值图层叠加后的统计值
    """
    # 打开文件
    raster = gdal.Open(input_raster)
    shapefile = ogr.Open(input_shapefile)
    layer = shapefile.GetLayer()

    # 获取栅格空间信息
    transform = raster.GetGeoTransform()
    x_origin = transform[0]
    y_origin = transform[3]
    pixel_width = transform[1]
    pixel_height = transform[5]

    # 将矢量数据空间坐标系投影到栅格
    source_sr = layer.GetSpatialRef()
    target_sr = osr.SpatialReference()
    target_sr.ImportFromWkt(raster.GetProjectionRef())
    coord_trans = osr.CoordinateTransformation(source_sr, target_sr)
    geom = feature.GetGeometryRef()
    geom.Transform(coord_trans)

    # 获取几何图形中的所有点
    if geom.GetGeometryName() == 'MULTIPOLYGON':
        count = 0
        points_x = []
        points_y = []

        for polygon in geom:
            geom_inner = geom.GetGeometryRef(count)
            ring = geom_inner.GetGeometryRef(0)
            numpoints = ring.GetPointCount()

            for p in range(numpoints):
                lon, lat, z = ring.GetPoint(p)
                points_x.append(lon)
                points_y.append(lat)

            count += 1
    elif geom.GetGeometryName() == 'POLYGON':
        ring = geom.GetGeometryRef(0)
        numpoints = ring.GetPointCount()
        points_x = []
        points_y = []

        for p in range(numpoints):
            lon, lat, z = ring.GetPoint(p)
            points_x.append(lon)
            points_y.append(lat)
    else:
        sys.exit("ERROR: Geometry needs to be either Polygon or Multipolygon")

    x_min = min(points_x)
    x_max = max(points_x)
    y_min = min(points_y)
    y_max = max(points_y)

    # 指定需要在栅格上读出的偏移及行列像素位
    x_offset = int((x_min - x_origin) / pixel_width)
    y_offset = int((y_origin - y_max) / pixel_width)
    x_count = int((x_max - x_min) / pixel_width) + 1
    y_count = int((y_max - y_min) / pixel_width) + 1

    # 在内存中创建一个栅格对象
    target_ds = gdal.GetDriverByName('MEM').Create('', x_count, y_count, 1, gdal.GDT_Byte)
    target_ds.SetGeoTransform((x_min, pixel_width, 0, y_max, 0, pixel_height))

    # 使对象栅格的投影坐标与原始栅格保持一致
    raster_srs = osr.SpatialReference()
    raster_srs.ImportFromWkt(raster.GetProjectionRef())
    target_ds.SetProjection(raster_srs.ExportToWkt())

    # 将区域多边形栅格化
    gdal.RasterizeLayer(target_ds, [1], layer, burn_values=[1])

    # 将栅格作为数组读入
    banddata_raster = raster.GetRasterBand(1)
    data_raster = banddata_raster.ReadAsArray(x_offset, y_offset, x_count, y_count).astype(np.float)

    band_mask = target_ds.GetRasterBand(1)
    data_mask = band_mask.ReadAsArray(0, 0, x_count, y_count).astype(np.float)

    # Mask zone of raster
    zone_raster = np.ma.masked_array(data_raster, np.logical_not(data_mask))

    # 计算当前区域的统计值
    return {'average': np.average(zone_raster), 'mean': np.mean(zone_raster),
            'median': np.ma.median(zone_raster), 'std': np.std(zone_raster),
            'var': np.var(zone_raster), 'histogram': str(np.histogram(zone_raster))}

def loop_zonal_stats(input_shapefile, input_raster):
    """遍历 shapefile 文件的所有图层
    分别将 shapefile 中的每个多边形与值图层叠加

    参数:
        input_shapefile (str): 输入 shapefile 文件路径
        input_raster (str): 输入 tif 文件路径

    返回:
        (array) 分区统计值
    """
    shapefile = ogr.Open(input_shapefile)
    layer = shapefile.GetLayer()
    features = range(layer.GetFeatureCount())
    stat_array = []

    for f_id in features:
        feature = layer.GetFeature(f_id)
        field_id = feature.GetField('id')
        f_value = zonal_stats(feature, input_shapefile, input_raster)
        f_value['id'] = field_id
        stat_array.append(f_value)
    return stat_array

# def raster_zonal_stats(zonal_raster, input_raster):
#     # 读取值图层文件
#     raster = gdal.Open(input_raster)
#     raster_band = raster.GetRasterBand(1)

#     # 获取值图层在当前坐标系的坐标对应关系
#     raster_trans = raster.GetGeoTransform()
#     x_origin = raster_trans[0]
#     y_origin = raster_trans[3]

#     # 读取分区图层及数据
#     region = gdal.Open(zonal_raster)
#     region_band = region.GetRasterBand(1)
#     region_array = region_band.ReadAsArray(0, 0, region_band.XSize, region_band.YSize)
#     # print(f"region no data value: {region_band.GetNoDataValue()}")

#     region_trans = region.GetGeoTransform()
#     region_x_origin = region_trans[0]
#     region_y_origin = region_trans[3]
#     region_pixel_width = region_trans[1]
#     region_pixel_height = region_trans[5]

#     # 计算相对偏移量
#     region_x_offset = (region_x_origin - x_origin) / region_pixel_width
#     region_y_offset = (y_origin - region_y_origin) / region_pixel_width

#     raster_array = raster_band.ReadAsArray(
#         region_x_offset,
#         region_y_offset,
#         region_band.XSize,
#         region_band.YSize
#     )

#     # print(f"raster_array: {raster_array}")

#     raster_zonal = np.ma.masked_array(
#         raster_array,
#         [x == region_band.GetNoDataValue() for x in region_array]
#     )

#     stat_dict = {}
#     stat_dict[0] = (np.average(raster_zonal), np.mean(raster_zonal), np.ma.median(raster_zonal),
#                     np.std(raster_zonal), np.var(raster_zonal), np.histogram(raster_zonal))
#     return stat_dict

def convert_tif_to_shp(tif_file, dst_layername='OUT'):
    """将 tif 文件转换成 shapefile 文件
    使用 gdal.Polygonize

    参数:
        tif_file (str): 输入 tif 文件路径
        dst_layername (str): 输出 shapefile 文件名

    返回:
        (str) 转换后的 shapefile 文件路径
    """
    src_ds = gdal.Open(tif_file)
    srcband = src_ds.GetRasterBand(1)

    srs = osr.SpatialReference()
    srs.ImportFromWkt(src_ds.GetProjection())

    drv = ogr.GetDriverByName("ESRI Shapefile")

    if os.path.exists(dst_layername+".shp"):
        drv.DeleteDataSource(dst_layername+".shp")

    dst_ds = drv.CreateDataSource(dst_layername + ".shp")
    dst_layer = dst_ds.CreateLayer(dst_layername, geom_type=ogr.wkbPolygon, srs=srs)

    df = ogr.FieldDefn('id', ogr.OFTInteger)
    dst_layer.CreateField(df)

    maskband = srcband.GetMaskBand()
    gdal.Polygonize(srcband, maskband, dst_layer, 0, [], callback=None)

    dst_ds.Destroy()
    src_ds = None

    return dst_layername + '.shp'

def raster_zonal_stats(zonal_raster, input_raster):
    """获取栅格数据的分区统计值
    将 tif 文件转成 shapefile, 将 shapefile 与值图层叠加计算

    参数:
        zonal_raster (str): 分区栅格图层路径
        input_raster (str): 值图层路径

    返回:
        (array) 分区统计值
    """
    shp_file = convert_tif_to_shp(zonal_raster)
    result = loop_zonal_stats(shp_file, input_raster)

    return result

def main(input_file, input_raster, options=None):
    """主函数
    依据分区图层的类型(矢量或栅格)，选择不同的处理方式。
    打印出 平均值, 算术平均值, 中位数, 标准差, 方差 和 直方图

    参数:
        input_file (str): 分区图层的路径
        input_raster(str): 值图层的路径
    """
    if input_file.endswith('.shp'):
        result = loop_zonal_stats(input_file, input_raster)
    elif input_file.endswith('.tif'):
        result = raster_zonal_stats(input_file, input_raster)
    else:
        print(f"invalid input file")
        sys.exit(1)

    if (options is None) or (options['format'] is None):
        for item in result:
            print(f"{item['id']} average: {item['average']} mean: {item['mean']} \
median: {item['median']} std: {item['std']} var: {item['var']} histogram: {item['histogram']}")

    elif options['format'] == 'csv':
        out_file = 'out.csv'
        if options['outFile'] is not None:
            out_file = options['outFile']

        with open(out_file, 'w', newline='') as csvfile:
            fieldnames = list(result[0].keys())
            writer = csv.DictWriter(csvfile, fieldnames=fieldnames)

            writer.writeheader()
            writer.writerows(result)

    elif options['format'] == 'json':
        out_file = 'out.json'
        if options['outFile'] is not None:
            out_file = options['outFile']

        with open(out_file, 'w') as jsonfile:
            json.dump(result, jsonfile, sort_keys=True, indent=4)

if __name__ == '__main__':
    argv = gdal.GeneralCmdLineProcessor(sys.argv)

    options = {}
    format = None
    out_filename = None
    src_filename = None
    dst_filename = None

    i = 1
    while i < len(argv):
        arg = argv[i]

        if arg == '-f' or arg == '-of':
            i = i + 1
            format = argv[i]

        elif arg == '-o' or arg == '-out':
            i = i + 1
            out_filename = argv[i]

        elif src_filename is None:
            src_filename = argv[i]

        elif dst_filename is None:
            dst_filename = argv[i]

        else:
            print("must have at least two arguments: zonal[.shp|.tif] raster.tif")
            sys.exit(1)

        i = i + 1

    options['format'] = format
    options['outFile'] = out_filename
    main(src_filename, dst_filename, options)
